#include "stdafx.h"
#include <map>

struct memmon_t
{
	struct mbdesc_t
	{
		size_t _ret;
		size_t _siz;
		mbdesc_t( ) : _ret(0), _siz(0)
		{}
		mbdesc_t( size_t ret, size_t bytes ) : _ret(ret), _siz(bytes)
		{
		}
	};
	
	typedef std::map<void*,mbdesc_t>	alloc_detail_map_t;

	struct _u64
	{
		uint64_t val;
		_u64() : val(0){}
		operator uint64_t& ()
		{
			return val;
		}
		operator const uint64_t () const 
		{
			return val;
		}
	};

	typedef std::map<size_t, _u64>	size_alloc_map_t;

	struct mod_info_t
	{
		std::string			modulename;
		size_t				alloctimes;
		size_t				freetimes;
		uint64_t			totalallocbytes;
		size_t				maxalloc;
		alloc_detail_map_t	detailmap;
		mod_info_t() 
			: alloctimes(0)
			, freetimes(0)
			, totalallocbytes(0)
			, maxalloc(0)
		{
		}
	};

	typedef std::map<LPCVOID,mod_info_t>	mod_log_map_t;

	mod_log_map_t		_modmap;
	size_alloc_map_t	_sizmap;

	memmon_t()
	{
		int a = 1;
	}

	~memmon_t()
	{
		mod_log_map_t::iterator it = _modmap.begin();
		for( ; it != _modmap.end(); ++ it )
		{
			mod_info_t& info = it->second;
			printf( "%s\n", info.modulename.c_str() );
			printf( "\talloctimes: %d\n", info.alloctimes );
			printf( "\tfreetimes: %d\n", info.freetimes );
			printf( "\ttotalallocbytes: %I64d\n", info.totalallocbytes );
			printf( "\tmaxalloc: %d\n", info.maxalloc );
			printf( "\tunfreed: %d\n", info.detailmap.size() );
			alloc_detail_map_t::iterator it2 = info.detailmap.begin();
			for( ; it2 != info.detailmap.end(); ++ it2 )
			{
				printf( "\t\tleakfrom: 0x%08x, %d(b)\n", it2->second._ret, it2->second._siz );
			}
		}
		size_alloc_map_t::iterator it2 = _sizmap.begin();
		for( ; it2 != _sizmap.end(); ++ it2 )
		{
			printf( "[%08x] : %I64d times\n", it2->first, (const uint64_t)(it2->second) );
		}

	}

	size_t get_rva(LPCVOID mb, void * ret )
	{
		return (size_t)ret - (size_t)mb;
	}

	void* logalloc( void * pmem, size_t cb, void * ret )
	{
		size_t acb = cb/16*16;
		_sizmap[acb] ++ ;

		MEMORY_BASIC_INFORMATION meminfo = {};
		VirtualQuery( ret, &meminfo, sizeof(meminfo) );
		LPCVOID hmodbase = meminfo.AllocationBase;
		mod_log_map_t::iterator it = _modmap.find(hmodbase);
		size_t urva = get_rva(hmodbase,ret);
		if( it == _modmap.end() )
		{
			mod_info_t mlog;
			char szname[512] = {};
			GetModuleFileNameA((HMODULE)hmodbase, szname, 512 );
			mlog.modulename = szname;
			mlog.alloctimes ++;
			mlog.totalallocbytes += cb;
			mlog.maxalloc = max(mlog.maxalloc,cb);
			mlog.detailmap[pmem] = mbdesc_t(urva,cb);
			_modmap[hmodbase] = mlog;
		}
		else
		{
			mod_info_t& mlog = it->second;
			mlog.alloctimes ++;
			mlog.totalallocbytes += cb;
			mlog.maxalloc = max(mlog.maxalloc,cb);
			mlog.detailmap[pmem] = mbdesc_t(get_rva(hmodbase,ret),cb);
		}
		return pmem;
	}
	void* logrealloc( void * po, void * pn, size_t cb, void * ret )
	{
		size_t acb = cb/16*16;
		_sizmap[acb] ++ ;
		MEMORY_BASIC_INFORMATION meminfo = {};
		VirtualQuery( ret,&meminfo, sizeof(meminfo) );
		LPCVOID hmodbase = meminfo.AllocationBase;
		mod_info_t& mlog = _modmap[hmodbase];
		mlog.alloctimes ++;
		mlog.freetimes ++;
		mlog.totalallocbytes += cb;
		mlog.maxalloc = max(mlog.maxalloc,cb);
		mlog.detailmap.erase(po);
		mlog.detailmap[pn] = mbdesc_t(get_rva(hmodbase,ret),cb);
		return pn;
	}
	void logfree( void* p, void * ret )
	{
		MEMORY_BASIC_INFORMATION meminfo = {};
		VirtualQuery( ret,&meminfo, sizeof(meminfo) );
		LPCVOID hmodbase =meminfo.AllocationBase;
		mod_info_t& mlog = _modmap[hmodbase];
		if( mlog.detailmap.find(p) == mlog.detailmap.end() )
			__asm int 3;
		mlog.freetimes ++;
		mlog.detailmap.erase(p);
	}

};


#define	get_ret_address(arg0)		(((void**)&arg0)[-1])

static memmon_t _leakmon;

static memmon_t & get_leakmon()
{
	return _leakmon;
}

void * __cdecl xm_malloc( size_t cb ) 
{
	if( !cb ) return NULL;
	void * p = malloc( cb );
	if( !p ) return p;
	return get_leakmon().logalloc( p, cb, get_ret_address(cb) );
}

void __cdecl xm_free( void * p ) 
{
	if( !p ) return ;
	free( p );
	get_leakmon().logfree(p, get_ret_address(p) );
}

void * __cdecl xm_realloc( void * p, size_t cb ) 
{ 
	void * px =  realloc( p, cb );
	if( !px ) return px;
	return get_leakmon().logrealloc( p, px, cb, get_ret_address(cb) );
}

void * __cdecl xm_calloc( size_t cnt, size_t size ) 
{ 
	void * p =  calloc( cnt, size );
	if( !p ) return p;
	return get_leakmon().logalloc( p, cnt*size, get_ret_address(cnt) );
}
